## de-bugged and ready to go! ##

############################### Functions #############################

initialise<-function(init.size, incidence){
	X<-runif(init.size, 0, 200)
	Y<-runif(init.size, 0, 100)
	parasite<-rbinom(init.size, 1, incidence)
	out<-cbind(X, Y, parasite)
	out
}
	
density<-function(popmatrix, cell.size, space.x, space.y){
	gridX<-1+popmatrix[,"X"]%/%cell.size	#converts to grid reference
	gridY<-1+popmatrix[,"Y"]%/%cell.size
	zeros<-rep(0, space.x/cell.size*space.y/cell.size)
	dens<-matrix(data=zeros, ncol=space.x/cell.size, nrow=space.y/cell.size)
	par.grid<-matrix(data=zeros, ncol=space.x/cell.size, nrow=space.y/cell.size)
	for (i in 1:length(gridX)){
		dens[gridY[i], gridX[i]]<-1+dens[gridY[i], gridX[i]]
		if (popmatrix[i, "parasite"]==1) par.grid[gridY[i], gridX[i]]<-1	#determine if parasites are present in each cell
	}
	local.dens<-rep(0, length(gridX))
	local.parasite<-c()
	for (i in 1:length(gridX)){
		local.dens[i]<-dens[gridY[i], gridX[i]]	
		par<-par.grid[popmatrix[i,"Y"]/cell.size+1, popmatrix[i,"X"]/cell.size+1]
		local.parasite<-c(local.parasite, par)
	}
	popmatrix<-cbind(popmatrix, local.dens, local.parasite) 
	popmatrix
}	

reproduce<-function(popmatrix, r, K, parasite.cost){
	#takes parasite load into account to calculate reproduction
	local.dens<-popmatrix[,"local.dens"]
	cost<-popmatrix[, "parasite"]+1
	cost[which(cost==2)]<-parasite.cost
	recruits<-(exp(r*(1-local.dens/K)))*cost
	recruits<-rpois(length(recruits), recruits)
	popmatrix<-cbind(popmatrix, recruits)
	popmatrix
}

disperse<-function(popmatrix, max.disp, space.x, space.y, incidence, parasite.cost){
	#disperses recruits based on normal cross-section kernel
	X<-c()
	Y<-c()
	min.x<-max(popmatrix[,"X"])-200
	if (min.x<0) min.x<-0
	parasite<-c()
	par.cost<-popmatrix[,"parasite"]+1
	par.cost[which(par.cost==2)]<-parasite.cost
	for (i in 1:length(popmatrix[,"X"])){
		dists<-rlnorm(popmatrix[i, "recruits"], log(5), log(2))
		psurv<-1-dists/max.disp #distance based survival
		psurv[which(psurv<0)]<-0
		psurv<-rbinom(length(psurv), 1, psurv)
		if (sum(psurv==0)>0) dists<-dists[-which(psurv==0)] else next
		dists<-dists*par.cost[i]
		parasitised<-rbinom(length(dists), 1, popmatrix[i, "local.parasite"]*incidence)
		directions<-runif(length(dists), 0, 2*pi) #random directions
		xchange<-dists*cos(directions)
		xchange<-xchange+popmatrix[i,"X"]
		ychange<-dists*sin(directions)
		ychange<-ychange+popmatrix[i, "Y"]
		XY<-reflect(xchange, ychange, space.x, space.y, min.x)
		X<-c(X, XY[,1])
		Y<-c(Y, XY[,2])
		parasite<-c(parasite, parasitised)
	}
	new.pop<-cbind(X, Y, parasite)
	new.pop
}

par.plot<-function(popmatrix){
	clean.x<-popmatrix[which(popmatrix[,"parasite"]==0), "X"]
	clean.y<-popmatrix[which(popmatrix[,"parasite"]==0), "Y"]
	inf.x<-popmatrix[which(popmatrix[,"parasite"]==1), "X"]
	inf.y<-popmatrix[which(popmatrix[,"parasite"]==1), "Y"]
	plot(c(clean.x, inf.x), c(clean.y, inf.y), type="n", xlab="x-axis distance (arbitrary units)", ylab="y-axis distance (arbitrary units)")
	points(clean.x, clean.y, pch=19, col="grey50")
	points(inf.x, inf.y, pch=19, col="black")
}

torus<-function(newx, newy, space.x, space.y){ #wraps space around a torus
	newx[which(newx>space.x)]<-newx[which(newx>space.x)]%%space.x #wrap space
	newx[which(newx<0)]<-space.x+newx[which(newx<0)]
	newy[which(newy>space.y)]<-newy[which(newy>space.y)]%%space.y #wrap space
	newy[which(newy<0)]<-space.y+newy[which(newy<0)]
	XY<-cbind(newx, newy)
	XY
}

reflect<-function(newx, newy, space.x, space.y, min.x){ #reflecting edges
	newx[which(newx>space.x)]<-2*space.x-newx[which(newx>space.x)]
	newx[which(newx<min.x)]<- 2*min.x-newx[which(newx<min.x)]
	newy[which(newy>space.y)]<-2*space.y-newy[which(newy>space.y)]
	newy[which(newy<0)]<- -newy[which(newy<0)]
	XY<-cbind(newx, newy)
	XY
}

rollup<-function(popmatrix){
	max.x<-max(popmatrix[,"X"])
	min.x<-max.x-200
	if (min.x<0) min.x<-0
	if (min(popmatrix[,"X"])<min.x) popmatrix<-popmatrix[-which(popmatrix[,"X"]<min.x),]
	popmatrix	
}

summ.init<-function(){ #initialises summary data matrix
	names<-c("disp.cost", "repro.cost", "surv.cost", "incidence", "r", "K", "rep", "x.diff", "free.count", "cell.size", "par.prop","edge.prop","edge.prop.n", "core.prop", "core.prop.n")	
	out<-matrix(data=rep(999, 11), nrow=1, ncol=15)
	colnames(out)<-names
	out
}

run.sampler<-function(popmatrix, runsumm, r, K, rep, disp.cost, repro.cost, incidence, surv.cost, space.x, cell.size){
	if (length(popmatrix[,"X"])>0){
		clean.max<-max(popmatrix[which(popmatrix[,"parasite"]==0),"X"])
		inf.max<-max(popmatrix[which(popmatrix[,"parasite"]==1), "X"])
		x.diff<-clean.max-inf.max
		free.count<-sum(popmatrix[,"X"]>inf.max)
		pop.max<-max(clean.max, inf.max)
		edge.prop<-sum(popmatrix[which(popmatrix[,"X"]>(pop.max-10)),"parasite"])/sum(popmatrix[,"X"]>(pop.max-10))
		edge.prop.n<-sum(popmatrix[,"X"]>(pop.max-10))
		core.prop<-sum(popmatrix[which(popmatrix[,"X"]<(pop.max-100)),"parasite"])/sum(popmatrix[,"X"]<(pop.max-100))
		core.prop.n<-sum(popmatrix[,"X"]<(pop.max-100))
		if (space.x-pop.max<10){
			x.diff<-NA
			free.count<-NA	
			edge.prop<-NA
			core.prop<-NA
		}
	par.prop<-sum(popmatrix[,"parasite"]==1)/length(popmatrix[,"parasite"])
		runsumm<-rbind(runsumm, c(disp.cost, repro.cost, surv.cost, incidence, r, K, rep, x.diff, free.count, cell.size, par.prop, edge.prop, edge.prop.n, core.prop, core.prop.n))
		print(c(incidence, cell.size, par.prop, rep))
	}
	runsumm
}
	
sim.equilibrium<-function(popmatrix, generations, space.x, space.y, r, K, cell.size, incidence, repro.cost, disp.cost, surv.cost, max.disp){
	for (g in 1:generations){
	if (length(popmatrix)>0){
		par.plot(popmatrix)
		popmatrix<-density(popmatrix, cell.size, space.x, space.y)
		popmatrix<-reproduce(popmatrix, r, K, parasite.cost=1-repro.cost)
		popmatrix<-disperse(popmatrix, max.disp, space.x, space.y, incidence, parasite.cost=1-disp.cost)
		#print(length(popmatrix[,"X"]))
	}
	}	
	popmatrix
}		
	
sim.expansion<-function(popmatrix, generations, space.x, space.y, r, K, cell.size, incidence, repro.cost, disp.cost, surv.cost, max.disp){
	for (g in 1:generations){
	if (length(popmatrix)>0){
		popmatrix<-rollup(popmatrix)
		par.plot(popmatrix)
		popmatrix<-density(popmatrix, cell.size, space.x, space.y)
		popmatrix<-reproduce(popmatrix, r, K, parasite.cost=1)
		popmatrix<-disperse(popmatrix, max.disp, space.x, space.y, incidence, parasite.cost=1)
		#print(length(popmatrix[,"X"]))	
	}
	}
	popmatrix
}	

############################# The Model ##############################
r<-1.5
repro.cost<-0
disp.cost<-0
surv.cost<-0
max.disp<-15
K<-20
cell.size<-5
KperArea<-K/cell.size^2
incidence<-0.3

popmatrix<-initialise(KperArea*2000, incidence)

popmatrix<-sim.equilibrium(popmatrix, 5, space.x=200, space.y=100, r, K, cell.size, incidence, repro.cost, disp.cost, surv.cost, max.disp)
			
popmatrix<-sim.expansion(popmatrix, 45, space.x=450, space.y=100, r, K, cell.size, incidence, repro.cost, disp.cost, surv.cost, max.disp)
			